#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
评估 JSONL 格式预测结果中 CWE 编号的提取和文本准确率。
用法示例：
    python evaluate_cwe.py -i predictions.jsonl [--verbose]
"""

import json
import re
import argparse
from pathlib import Path
from typing import Optional, List, Dict
import os

# 正则只提取数字部分，返回形如 "CWE-123" 的字符串
CWE_PATTERN = re.compile(r'CWE-(\d+)', re.IGNORECASE)

def extract_cwe(text: str) -> Optional[str]:
    """
    从文本中提取第一个 CWE-<数字>，返回标准格式的字符串；未匹配时返回 None。
    """
    match = CWE_PATTERN.search(text or "")
    return f"CWE-{match.group(1)}" if match else None

def load_records(path: Path) -> List[Dict[str, str]]:
    """
    从 JSONL 文件读取每行记录，返回列表。忽略解析出错的行并打印警告。
    """
    records = []
    with path.open('r', encoding='utf-8') as f:
        for lineno, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                records.append(json.loads(line))
            except json.JSONDecodeError as e:
                print(f"[Warning] 第 {lineno} 行 JSON 解析失败: {e}")
    return records

def evaluate(records: List[Dict[str, str]], output_file: Optional[Path] = None, verbose: bool = False) -> None:
    """
    计算并打印：
      - 样本总数
      - CWE 编号准确率（预测编号在标签编号列表中即视为正确）
      - 原始文本完全匹配准确率
    verbose 模式下，会打印每个编号不匹配的详细信息。
    如果指定了 output_file，则将结果写入文件。
    """
    total = len(records)
    correct_cwe = 0
    text_correct = 0

    for idx, rec in enumerate(records, 1):
        pred_txt = rec.get('predict', '')
        label_txt = rec.get('label', '')

        cw_pred = extract_cwe(pred_txt)
        # 解析 label 字段：JSON 列表或纯文本
        true_cs: List[str] = []
        if label_txt.startswith('[') and label_txt.endswith(']'):
            try:
                items = json.loads(label_txt)
            except json.JSONDecodeError:
                items = [label_txt]
        else:
            items = [label_txt]
        for it in items:
            cw = extract_cwe(str(it))
            if cw:
                true_cs.append(cw)
        if not true_cs:
            true_cs = [""]  # 没提取到任何编号时用空串占位

        # 统计 CWE 编号正确数
        if cw_pred in true_cs:
            correct_cwe += 1

        # 文本级别匹配
        if pred_txt.strip() == label_txt.strip():
            text_correct += 1
        elif verbose:
            print(f"[Mismatch #{idx}] predict=\"{pred_txt}\" | label=\"{label_txt}\" | "
                  f"cw_pred={cw_pred} | true_cs={true_cs}")

    cwe_acc = correct_cwe / total if total else 0.0
    text_acc = text_correct / total if total else 0.0

    result = (
        f"样本总数：{total}\n"
        f"CWE 编号准确率：{cwe_acc:.2%} （{correct_cwe}/{total}）\n"
        f"原始文本准确率：{text_acc:.2%} （{text_correct}/{total}）\n"
    )

    print(result)

    if output_file:
        output_file = Path(output_file)  # Ensure output_file is a Path object
        with output_file.open('w', encoding='utf-8') as f:
            f.write(result)

def main():
    parser = argparse.ArgumentParser(description="计算 JSONL 中 CWE 编号预测准确率")
    parser.add_argument('-i', '--input', type=Path, required=True,
                        help='输入 JSONL 文件，需包含 predict 和 label 字段')
    parser.add_argument('--verbose', action='store_true', default=False,
                        help='打印每条不匹配的详细信息')
    parser.add_argument('-o', '--output', type=Path, default=None,
                        help='输出结果到指定文件，默认为标准输出')
    args = parser.parse_args()

    if not args.input.is_file():
        print(f"[Error] 找不到文件：{args.input}")
        return

    records = load_records(args.input)
    if not records:
        print(f"[Error] 未从 {args.input} 读取到任何记录。")
        return

    # 简单字段检查
    for fld in ('predict', 'label'):
        if fld not in records[0]:
            print(f"[Error] 记录中缺少字段：'{fld}'")
            return
    if args.output:
        outputpath = args.output
    else:
        current_dir = os.path.dirname(os.path.abspath(__file__))
        ## 获取input文件名前缀
        input_filename = args.input.stem
        ## 生成输出文件路径
        outputpath = os.path.join(current_dir, 'res',f'{input_filename}_eval.txt')


    evaluate(records, output_file=outputpath, verbose=args.verbose)

if __name__ == '__main__':
    main()
